"""
音频特征提取模块
功能：从语音文件中提取MFCC、ZCR、Chroma、Energy、Pitch等传统声学特征
参数：音频文件路径或音频数据数组
返回值：提取的特征向量（numpy数组）
使用场景：语音情感识别的特征预处理步骤
"""

import librosa
import numpy as np
import scipy.stats
from typing import Dict, Any, Tuple
from core.config import settings

class AudioFeatureExtractor:
    """
    音频特征提取器类
    用途：提取语音文件的各种传统声学特征
    """
    
    def __init__(self):
        """初始化特征提取器"""
        self.sample_rate = settings.SAMPLE_RATE
        self.hop_length = settings.HOP_LENGTH
        self.mfcc_components = settings.MFCC_COMPONENTS
        
    def load_audio(self, file_path: str) -> Tuple[np.ndarray, int]:
        """
        加载音频文件
        参数：file_path - 音频文件路径
        返回值：音频数据数组和采样率
        """
        try:
            # 加载音频，限制时长为指定秒数
            audio_data, sr = librosa.load(
                file_path, 
                sr=self.sample_rate, 
                duration=settings.AUDIO_DURATION
            )
            return audio_data, sr
        except Exception as e:
            raise ValueError(f"音频文件加载失败: {str(e)}")
    
    def extract_mfcc_features(self, audio_data: np.ndarray) -> Dict[str, float]:
        """
        提取MFCC特征
        参数：audio_data - 音频数据数组
        返回值：MFCC统计特征字典
        """
        # 提取MFCC特征
        mfcc = librosa.feature.mfcc(
            y=audio_data,
            sr=self.sample_rate,
            n_mfcc=self.mfcc_components,
            hop_length=self.hop_length
        )
        
        # 计算统计特征
        features = {}
        features['mfcc_mean'] = np.mean(mfcc, axis=1)
        features['mfcc_std'] = np.std(mfcc, axis=1)
        features['mfcc_delta'] = np.mean(librosa.feature.delta(mfcc), axis=1)
        
        return features
    
    def extract_zcr_features(self, audio_data: np.ndarray) -> Dict[str, float]:
        """
        提取过零率(ZCR)特征
        参数：audio_data - 音频数据数组
        返回值：ZCR统计特征字典
        """
        zcr = librosa.feature.zero_crossing_rate(
            audio_data, 
            hop_length=self.hop_length
        )[0]
        
        return {
            'zcr_mean': np.mean(zcr),
            'zcr_std': np.std(zcr),
            'zcr_median': np.median(zcr)
        }
    
    def extract_chroma_features(self, audio_data: np.ndarray) -> Dict[str, float]:
        """
        提取色度图谱(Chroma)特征
        参数：audio_data - 音频数据数组
        返回值：Chroma统计特征字典
        """
        chroma = librosa.feature.chroma_stft(
            y=audio_data,
            sr=self.sample_rate,
            hop_length=self.hop_length
        )
        
        return {
            'chroma_mean': np.mean(chroma, axis=1),
            'chroma_std': np.std(chroma, axis=1),
            'chroma_var': np.var(chroma, axis=1)
        }
    
    def extract_energy_features(self, audio_data: np.ndarray) -> Dict[str, float]:
        """
        提取能量特征
        参数：audio_data - 音频数据数组
        返回值：能量统计特征字典
        """
        # 计算RMS能量
        rms_energy = librosa.feature.rms(
            y=audio_data,
            hop_length=self.hop_length
        )[0]
        
        # 计算谱质心
        spectral_centroid = librosa.feature.spectral_centroid(
            y=audio_data,
            sr=self.sample_rate,
            hop_length=self.hop_length
        )[0]
        
        return {
            'rms_mean': np.mean(rms_energy),
            'rms_std': np.std(rms_energy),
            'spectral_centroid_mean': np.mean(spectral_centroid),
            'spectral_centroid_std': np.std(spectral_centroid),
            'energy_variance': np.var(rms_energy)
        }
    
    def extract_pitch_features(self, audio_data: np.ndarray) -> Dict[str, float]:
        """
        提取基频(Pitch)特征
        参数：audio_data - 音频数据数组
        返回值：基频统计特征字典
        """
        # 提取基频
        pitches, magnitudes = librosa.piptrack(
            y=audio_data,
            sr=self.sample_rate,
            hop_length=self.hop_length
        )
        
        # 获取每帧最强的基频
        pitch_values = []
        for t in range(pitches.shape[1]):
            index = magnitudes[:, t].argmax()
            pitch = pitches[index, t]
            if pitch > 0:
                pitch_values.append(pitch)
        
        if len(pitch_values) == 0:
            return {
                'pitch_mean': 0.0,
                'pitch_std': 0.0,
                'pitch_median': 0.0,
                'pitch_range': 0.0
            }
        
        pitch_array = np.array(pitch_values)
        return {
            'pitch_mean': np.mean(pitch_array),
            'pitch_std': np.std(pitch_array),
            'pitch_median': np.median(pitch_array),
            'pitch_range': np.max(pitch_array) - np.min(pitch_array)
        }
    
    def extract_all_features(self, file_path: str) -> np.ndarray:
        """
        提取所有特征（兼容训练模型的40维特征）
        参数：file_path - 音频文件路径
        返回值：合并后的特征向量
        使用场景：情感识别模型的输入特征
        """
        # 加载音频
        audio_data, _ = self.load_audio(file_path)
        
        # 提取各类特征
        mfcc_features = self.extract_mfcc_features(audio_data)
        zcr_features = self.extract_zcr_features(audio_data)
        chroma_features = self.extract_chroma_features(audio_data)
        energy_features = self.extract_energy_features(audio_data)
        pitch_features = self.extract_pitch_features(audio_data)
        
        # 合并特征，确保输出40维特征
        feature_vector = []
        
        # MFCC特征：使用前10个组件的均值和标准差 (20维)
        mfcc_mean = mfcc_features['mfcc_mean'][:10]  # 前10个MFCC系数
        mfcc_std = mfcc_features['mfcc_std'][:10]    # 前10个MFCC系数的标准差
        feature_vector.extend(mfcc_mean)
        feature_vector.extend(mfcc_std)
        
        # ZCR特征：3维
        feature_vector.extend(zcr_features.values())
        
        # Energy特征：5维
        feature_vector.extend(energy_features.values())
        
        # Pitch特征：4维
        feature_vector.extend(pitch_features.values())
        
        # Chroma特征：使用前8个色度的均值 (8维)
        chroma_mean = chroma_features['chroma_mean'][:8]
        feature_vector.extend(chroma_mean)
        
        # 确保特征向量长度为40
        feature_array = np.array(feature_vector)
        if len(feature_array) != 40:
            # 如果长度不对，截断或填充到40维
            if len(feature_array) > 40:
                feature_array = feature_array[:40]
            else:
                # 用零填充到40维
                padding = np.zeros(40 - len(feature_array))
                feature_array = np.concatenate([feature_array, padding])
        
        return feature_array
    
    def get_config(self) -> Dict[str, Any]:
        """
        获取特征提取器配置
        返回值：配置信息字典
        """
        return {
            'sample_rate': self.sample_rate,
            'hop_length': self.hop_length,
            'mfcc_components': self.mfcc_components,
            'feature_dimension': 40,
            'audio_duration': settings.AUDIO_DURATION,
            'features': [
                'mfcc_mean', 'mfcc_std', 'zcr', 'energy', 'pitch', 'chroma'
            ]
        } 